﻿using System;
using System.Buffers;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using PooledGrowableBufferHelper;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using WholeSlideImageLibrary.DeepZoom;

namespace WholeSlideImageLibrary.Formats.OptimizedPyramidTiledImage
{
    internal class OptimizedPyramidTiledImageDeepZoomTileSource : DeepZoomTileSource
    {
        private readonly OptimizedPyramidTiledWholeSlideImage _image;

        private LayerDecoder[]? _layers;

        public OptimizedPyramidTiledImageDeepZoomTileSource(OptimizedPyramidTiledWholeSlideImage image)
        {
            _image = image;

            Initialize();
        }

        public override EncodedImageFileFormat TileFileFormat => EncodedImageFileFormat.Jpeg;

        public override string GetDefinitionFile()
        {
            return "<?xml version=\"1.0\" encoding=\"UTF-8\"?><Image xmlns=\"http://schemas.microsoft.com/deepzoom/2008\" Format=\"jpg\" Overlap=\"0\" TileSize=\"" + _image.TileSize + "\"><Size Height=\"" + _image.Height + "\" Width=\"" + _image.Width + "\" /></Image>";
        }

        public override bool CheckTileExists(int layer, int column, int row)
        {
            LayerDecoder[]? layers = _layers;
            if (layers is null)
            {
                throw new InvalidOperationException(nameof(OptimizedPyramidTiledImageDeepZoomTileSource) + " is not properly initialized.");
            }
            if ((uint)layer >= (uint)layers.Length)
            {
                return false;
            }

            return layers[layer].CheckTileExists(row, column);
        }

        public override Task WriteToBufferAsync(int layer, int column, int row, IBufferWriter<byte> writer, CancellationToken cancellationToken)
        {
            LayerDecoder[]? layers = _layers;
            if (layers is null)
            {
                throw new InvalidOperationException(nameof(OptimizedPyramidTiledImageDeepZoomTileSource) + " is not properly initialized.");
            }
            if ((uint)layer >= (uint)layers.Length)
            {
                throw new ArgumentOutOfRangeException(nameof(layer));
            }

            return layers[layer].WriteTileAsync(row, column, writer, cancellationToken);
        }

        public override Task WriteToStreamAsync(int layer, int column, int row, Stream stream, CancellationToken cancellationToken)
        {
            LayerDecoder[]? layers = _layers;
            if (layers is null)
            {
                throw new InvalidOperationException(nameof(OptimizedPyramidTiledImageDeepZoomTileSource) + " is not properly initialized.");
            }
            if ((uint)layer >= (uint)layers.Length)
            {
                throw new ArgumentOutOfRangeException(nameof(layer));
            }

            return layers[layer].WriteTileAsync(row, column, stream, cancellationToken);
        }

        private void Initialize()
        {
            OptimizedPyramidTiledWholeSlideImage? image = _image;
            int levelCount = image.LevelCount;

            // 构建Level信息
            var levels = new List<LevelInformation>(levelCount);
            System.Drawing.Size size = image.GetLevelSize(0);
            var level0 = new LevelInformation(0, size.Width, size.Height, downsample: 1);
            levels.Add(level0);
            for (int i = 1; i < levelCount; i++)
            {
                size = image.GetLevelSize(i);
                double downsample = Math.Max(level0.Width / (double)size.Width, level0.Height / (double)size.Height);
                levels.Add(new LevelInformation(i, size.Width, size.Height, downsample));
            }

            // 构建Layers数组
            DeepZoomLayer[] dzLayers = DeepZoomHelper.CalculateDeepZoomLayers(level0.Width, level0.Height);
            LayerDecoder[] layers = new LayerDecoder[dzLayers.Length];

            for (int i = 0; i < layers.Length; i++)
            {
                DeepZoomLayer dzLayer = dzLayers[i];
                int layerDownsample = 1 << (dzLayers.Length - i - 1);

                // 找到合适的Level
                LevelInformation? suitableLevel = null;
                if (layerDownsample <= level0.Downsample)
                {
                    suitableLevel = level0;
                }
                if (suitableLevel is null)
                {
                    for (int j = 1; j < levels.Count; j++)
                    {
                        if (layerDownsample < levels[j].Downsample)
                        {
                            suitableLevel = levels[j - 1];
                            break;
                        }
                    }
                }
                if (suitableLevel is null)
                {
                    suitableLevel = levels[levels.Count - 1];
                }

                // 选择Decoder
                LevelInformation level = suitableLevel.GetValueOrDefault();
                if (dzLayer.Width == level.Width && dzLayer.Height == level.Height)
                {
                    layers[i] = new FastLayerDecoder(image)
                    {
                        Level = level.Index,
                        LayerWidth = dzLayer.Width,
                        LayerHeight = dzLayer.Height
                    };
                }
                else
                {
                    layers[i] = new ScaledLayerDecoder(image)
                    {
                        Level = level.Index,
                        LayerWidth = dzLayer.Width,
                        LayerHeight = dzLayer.Height,
                        LayerDownsample = layerDownsample,
                        LevelDownsample = level.Downsample
                    };
                }
            }

            _layers = layers;
        }

        readonly struct LevelInformation
        {
            public LevelInformation(int index, int width, int height, double downsample)
            {
                Index = index;
                Width = width;
                Height = height;
                Downsample = downsample;
            }

            public int Index { get; }
            public int Width { get; }
            public int Height { get; }
            public double Downsample { get; }
        }

        abstract class LayerDecoder
        {
            public LayerDecoder(OptimizedPyramidTiledWholeSlideImage image)
            {
                Image = image;
            }

            public OptimizedPyramidTiledWholeSlideImage Image { get; }

            public int Level { get; set; }
            public int LayerWidth { get; set; }
            public int LayerHeight { get; set; }

            public int TileSize => Image.TileSize;
            protected const int Overlap = 0;

            public bool CheckTileExists(int rowIndex, int columnIndex)
            {
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    return false;
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    return false;
                }
                return true;
            }

            public abstract Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken);

            public abstract Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken);
        }

        class FastLayerDecoder : LayerDecoder
        {
            public FastLayerDecoder(OptimizedPyramidTiledWholeSlideImage image) : base(image) { }

            private const int MaximumSegmentSize = 16384;

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken)
            {
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(columnIndex));
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(rowIndex));
                }

                OptimizedPyramidTiledWholeSlideImage? image = Image;
                FileSourceReaderCache? cache = Image.GetReaderCache();
                FileSourceReader? reader = null;
                byte[]? buffer = null;
                try
                {
                    reader = await cache.RentAsync(cancellationToken).ConfigureAwait(false);

                    FileContentRegion region = await image.GetTileStreamPositionAsync(reader, Level, columnIndex, rowIndex, cancellationToken).ConfigureAwait(false);
                    if (region.IsEmpty)
                    {
                        throw new InvalidDataException();
                    }

                    long offset = region.Offset;
                    int length = region.Length;

                    buffer = ArrayPool<byte>.Shared.Rent(Math.Min(length, MaximumSegmentSize));

                    while (length > 0)
                    {
                        int bytesToRead = Math.Min(length, MaximumSegmentSize);
                        int bytesRead = await reader.ReadAsync(offset, buffer.AsMemory(0, bytesToRead), cancellationToken).ConfigureAwait(false);
                        if (bytesToRead != bytesRead)
                        {
                            throw new InvalidDataException();
                        }

                        await stream.WriteAsync(buffer, 0, bytesToRead, cancellationToken).ConfigureAwait(false);

                        offset += bytesToRead;
                        length -= bytesToRead;
                    }
                }
                finally
                {
                    if (!(buffer is null))
                    {
                        ArrayPool<byte>.Shared.Return(buffer);
                    }
                    if (!(reader is null))
                    {
                        await cache.ReturnAsync(reader).ConfigureAwait(false);
                    }
                }
            }

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(columnIndex));
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(rowIndex));
                }

                OptimizedPyramidTiledWholeSlideImage? image = Image;
                FileSourceReaderCache? cache = Image.GetReaderCache();
                FileSourceReader? reader = null;
                try
                {
                    reader = await cache.RentAsync(cancellationToken).ConfigureAwait(false);

                    FileContentRegion region = await image.GetTileStreamPositionAsync(reader, Level, columnIndex, rowIndex, cancellationToken).ConfigureAwait(false);
                    if (region.IsEmpty)
                    {
                        throw new InvalidDataException();
                    }

                    long offset = region.Offset;
                    int length = region.Length;

                    while (length > 0)
                    {
                        int bytesToRead = Math.Min(length, MaximumSegmentSize);
                        Memory<byte> buffer = writer.GetMemory(bytesToRead);
                        bytesToRead = Math.Min(buffer.Length, length);

                        int bytesRead = await reader.ReadAsync(offset, buffer.Slice(0, bytesToRead), cancellationToken).ConfigureAwait(false);
                        if (bytesToRead != bytesRead)
                        {
                            throw new InvalidDataException();
                        }

                        writer.Advance(bytesRead);

                        offset += bytesToRead;
                        length -= bytesToRead;
                    }
                }
                finally
                {
                    if (!(reader is null))
                    {
                        await cache.ReturnAsync(reader).ConfigureAwait(false);
                    }
                }
            }
        }

        class ScaledLayerDecoder : LayerDecoder
        {
            public ScaledLayerDecoder(OptimizedPyramidTiledWholeSlideImage image) : base(image) { }

            public double LevelDownsample { get; set; }
            public int LayerDownsample { get; set; }

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, Stream stream, CancellationToken cancellationToken)
            {
                // 检查参数
                int tileSize = TileSize;
                int horizontalTileCount = (LayerWidth + tileSize - 1) / tileSize;
                int verticalTileCount = (LayerHeight + tileSize - 1) / tileSize;
                if ((uint)columnIndex >= (uint)horizontalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(rowIndex));
                }
                if ((uint)rowIndex >= (uint)verticalTileCount)
                {
                    throw new ArgumentOutOfRangeException(nameof(columnIndex));
                }

                int layerXOffset = columnIndex * tileSize;
                int layerYOffset = rowIndex * tileSize;
                int layerReadWidth = Math.Min(tileSize, LayerWidth - layerXOffset);
                int layerReadHeight = Math.Min(tileSize, LayerHeight - layerYOffset);

                // 转换成Level参数
                double downsample = LayerDownsample / LevelDownsample;
                int levelReadWidth = (int)(layerReadWidth * downsample);
                int levelReadHeight = (int)(layerReadHeight * downsample);
                int levelXOffset = (int)(layerXOffset * downsample);
                int levelYOffset = (int)(layerYOffset * downsample);

                // 读取并缩放
                Image<Rgba32>? scaledTile = null;
                byte[] buffer = ArrayPool<byte>.Shared.Rent(4 * levelReadWidth * levelReadHeight);
                try
                {
                    Memory<Rgba32> memory = buffer.AsMemory(0, 4 * levelReadWidth * levelReadHeight).Cast<byte, Rgba32>();
                    await Image.ReadRegionAsync(Level, levelXOffset, levelYOffset, levelReadWidth, levelReadHeight, memory, cancellationToken).ConfigureAwait(false);
                    using var tile = SixLabors.ImageSharp.Image.WrapMemory(memory, levelReadWidth, levelReadHeight);
                    scaledTile = tile.Clone(ctx => ctx.Resize(layerReadWidth, layerReadHeight));
                    scaledTile.SaveAsJpeg(stream);
                }
                finally
                {
                    ArrayPool<byte>.Shared.Return(buffer);
                    scaledTile?.Dispose();
                }
            }

            public override async Task WriteTileAsync(int rowIndex, int columnIndex, IBufferWriter<byte> writer, CancellationToken cancellationToken)
            {
                using (PooledMemoryStream? stream = PooledMemoryStreamManager.Shared.GetStream())
                {
                    await WriteTileAsync(rowIndex, columnIndex, (Stream)stream, cancellationToken).ConfigureAwait(false);

                    Copy(stream.ToReadOnlySequence(), writer);
                }
            }

        }

        private static void Copy(ReadOnlySequence<byte> source, IBufferWriter<byte> destination)
        {
            long length = source.Length;
            while (length > 0)
            {
                Span<byte> span = destination.GetSpan();
                int copyLength = (int)Math.Min(length, span.Length);
                source.Slice(0, copyLength).CopyTo(span);
                source = source.Slice(copyLength);
                destination.Advance(copyLength);
                length -= copyLength;
            }
        }
    }
}
